import streamlit as st
import asyncio
from loguru import logger
from typing import Dict, List, Any
import json
import random
import sys
import atexit
from pathlib import Path
import time
import re  # 导入正则表达式库用于分割文本

# 添加当前目录到搜索路径
sys.path.append(str(Path(__file__).parent))

# 使用相对导入
from utils import get_tools_from_mcp, call_tongyi_api, call_tongyi_api_with_messages, format_tool_call, format_tool_result, execute_tool_call, safe_disconnect, check_connection_health
from dify_utils import DifyKnowledgeBase
from config import DEFAULT_SSE_HOST, DEFAULT_SSE_PORT, DEFAULT_DATASET_ID

# 页面配置
st.set_page_config(
    page_title="AI 智能导购助手",
    page_icon="🛍️",
    layout="wide"
)

# 自定义CSS样式
st.markdown("""
    <style>
        /* 整体容器样式 */
        .block-container {
            padding: 2rem 3rem;
            max-width: 1000px;
            margin: 0 auto;
            background-color: #ffffff;
        }
        
        /* 标题样式 */
        .main .block-container h1 {
            color: #1e1e1e;
            font-size: 1.5rem;
            font-weight: 500;
            margin-bottom: 1.5rem;
        }
        
        /* 按钮基础样式 */
        div[data-testid="stButton"] button {
            width: 100%;
            text-align: left;
            padding: 0.75rem 1rem;
            background: #f7f9fc;
            border: 1px solid #d0d5dd;
            border-radius: 8px;
            color: #1e1e1e;
            font-size: 0.875rem;
            line-height: 1.5;
            font-weight: normal;
            margin: 0.25rem 0;
            transition: all 0.2s ease;
        }
        
        /* 按钮悬停效果 */
        div[data-testid="stButton"] button:hover {
            background: #f0f2f6;
            border-color: #98a2b3;
        }
        
        /* 聊天消息容器 */
        .stChatMessage {
            background: #ffffff;
            border: 1px solid #eaecf0;
            border-radius: 8px;
            padding: 1rem;
            margin: 0.5rem 0;
            box-shadow: 0 1px 2px rgba(16, 24, 40, 0.05);
        }
        
        /* 用户消息样式 */
        .stChatMessage [data-testid="StyledLinkIconContainer"] {
            background: #f7f9fc;
        }
        
        /* 系统消息样式 */
        .system-message {
            background: #f7f9fc;
            border: 1px solid #eaecf0;
            border-radius: 8px;
            padding: 0.75rem;
            margin: 0.5rem 0;
            font-size: 0.875rem;
            color: #475467;
        }
        
        /* 输入框样式 */
        .stTextInput input {
            border: 1px solid #d0d5dd;
            border-radius: 8px;
            padding: 0.75rem;
            font-size: 0.875rem;
            background: #ffffff;
        }
        
        .stTextInput input:focus {
            border-color: #98a2b3;
            box-shadow: 0 0 0 3px rgba(16, 24, 40, 0.05);
        }
        
        /* 列布局样式 */
        .row-widget.stHorizontal {
            gap: 1rem;
        }
        
        /* 移动端适配 */
        @media (max-width: 768px) {
            .block-container {
                padding: 1rem;
            }
            
            div[data-testid="stButton"] button {
                padding: 0.625rem 0.875rem;
            }
        }
    </style>
""", unsafe_allow_html=True)

# 示例查询（联想电脑相关）
EXAMPLE_QUERIES = [
    "我想购买联想ThinkPad X1 Carbon，预算10000元",
    "推荐几款联想品牌的高性能游戏本",
    "查找联想Yoga系列的轻薄本，价格不超过8000元",
    "了解联想拯救者系列的游戏本配置和价格",
    "询问联想笔记本电脑的最新优惠活动",
    "寻找联想商用笔记本的推荐型号",
    "查询联想平板电脑的相关信息",
    "获取联想笔记本电脑的售后服务政策",
]

# 初始化会话状态
if "messages" not in st.session_state:
    st.session_state.messages = []

if "tools" not in st.session_state:
    st.session_state.tools = []

if "sse_url" not in st.session_state:
    st.session_state.sse_url = "http://127.0.0.1:8100/sse"

if "is_connected" not in st.session_state:
    st.session_state.is_connected = False

if "selected_dataset" not in st.session_state:
    st.session_state.selected_dataset = DEFAULT_DATASET_ID

if "datasets" not in st.session_state:
    st.session_state.datasets = []

if "use_knowledge_base" not in st.session_state:
    st.session_state.use_knowledge_base = False

# 注册应用退出时的清理函数
def cleanup_on_exit():
    logger.info("应用退出，清理资源...")
    # 执行异步函数的包装器
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        loop.run_until_complete(safe_disconnect())
    finally:
        loop.close()
    logger.info("资源清理完成")

atexit.register(cleanup_on_exit)

# 处理服务器连接的函数
def connect_to_server():
    try:
        with st.spinner("正在连接服务器..."):
            # 获取工具列表
            tools = asyncio.run(get_tools_from_mcp(st.session_state.sse_url))
            if tools:
                st.session_state.tools = tools
                st.session_state.is_connected = True
                st.success("✅ 服务器连接成功！")
                logger.info(f"成功连接到 SSE 服务器: {st.session_state.sse_url}")
                return True
            else:
                st.session_state.is_connected = False
                st.error("❌ 获取工具列表失败，请检查服务器配置")
                logger.error(f"获取工具列表失败: {st.session_state.sse_url}")
                return False
    except Exception as e:
        st.session_state.is_connected = False
        st.error(f"❌ 连接服务器失败: {str(e)}")
        logger.error(f"连接服务器失败: {str(e)}")
        return False

# 加载知识库列表
def load_datasets():
    try:
        datasets = DifyKnowledgeBase.get_datasets()
        if datasets:
            st.session_state.datasets = datasets
            logger.info(f"成功加载知识库列表: {len(datasets)} 个知识库")
        else:
            logger.warning("未找到可用的知识库")
    except Exception as e:
        logger.error(f"加载知识库列表失败: {str(e)}")

# 检查连接健康状态的函数
def check_connection():
    with st.spinner("正在检查连接状态..."):
        is_healthy = asyncio.run(check_connection_health())
        if not is_healthy:
            st.session_state.is_connected = False
            st.warning("⚠️ 连接已断开，请重新连接")
            logger.warning("连接健康检查失败，已标记为断开")
            return False
        return True

# 处理用户查询的函数
async def process_user_query(prompt: str):
    # 记录开始时间
    start_time = time.time()
    
    # 显示处理中的消息
    with st.chat_message("assistant"):
        thinking = st.empty()
        thinking.markdown("🤔 思考中...")
    
    try:
        # 1. 如果启用了知识库，先查询知识库
        original_prompt = prompt
        knowledge_text = ""
        if st.session_state.use_knowledge_base:
            with st.chat_message("system"):
                st.info("🔍 正在从知识库检索相关内容...")
            
            knowledge_result = DifyKnowledgeBase.retrieve_knowledge(
                st.session_state.selected_dataset,
                prompt
            )
            
            # 如果有知识库内容，将其添加到提示中
            if knowledge_result:
                knowledge_text = DifyKnowledgeBase.format_knowledge_result(knowledge_result)
                if knowledge_text:
                    # 显示知识库检索结果
                    with st.chat_message("system"):
                        st.success("✅ 找到相关知识")
                        with st.expander("查看知识库内容", expanded=False):
                            st.markdown(knowledge_text)
                    prompt = f"{knowledge_text}\n\n用户问题：{prompt}"
                else:
                    with st.chat_message("system"):
                        st.warning("⚠️ 未找到相关知识")
            else:
                with st.chat_message("system"):
                    st.warning("⚠️ 知识库检索失败")
        
        # 2. 调用通义千问 API
        response = call_tongyi_api(prompt, st.session_state.tools)
        
        # 清除思考中的消息
        thinking.empty()
        
        # 计算总耗时
        total_time = time.time() - start_time
        
        # 3. 处理 API 响应
        if "error" in response:
            error_message = f"API 调用失败: {response['error']}"
            st.session_state.messages.append({"role": "assistant", "content": error_message})
            with st.chat_message("assistant"):
                st.error(error_message)
            return
        
        # 检查API响应格式
        output = response.get("output", {})
        if not output:
            error_message = "API 返回了空响应"
            st.session_state.messages.append({"role": "assistant", "content": error_message})
            with st.chat_message("assistant"):
                st.error(error_message)
            return
            
        # 处理新的API响应格式
        if "choices" in output:
            # 获取助手回复
            choice = output["choices"][0]
            message = choice.get("message", {})
            assistant_message = message.get("content", "")
            
            # 显示助手回复
            if assistant_message and not message.get("tool_calls"):
                # 如果使用了知识库，在回复前添加标记
                if knowledge_text:
                    assistant_message = f"📚 基于知识库内容回答：\n\n{assistant_message}"
                st.session_state.messages.append({"role": "assistant", "content": assistant_message})
                with st.chat_message("assistant"):
                    # 直接显示回复内容，不使用打字效果
                    st.markdown(assistant_message)
            
            # 处理工具调用
            tool_calls = message.get("tool_calls", [])
            if tool_calls:
                # 用于构建后续发送给模型的上下文
                tool_messages = []
                
                # 显示工具调用开始
                with st.chat_message("system"):
                    st.info("🛠️ 开始执行工具调用...")
                
                for tool_call in tool_calls:
                    tool_name = tool_call.get("function", {}).get("name", "")
                    arguments_str = tool_call.get("function", {}).get("arguments", "{}")
                    
                    # 解析参数
                    try:
                        arguments = json.loads(arguments_str)
                    except:
                        arguments = {}
                    
                    # 显示工具调用信息
                    with st.chat_message("system"):
                        st.info(f"📝 调用工具：{tool_name}")
                        with st.expander("查看调用参数", expanded=False):
                            st.json(arguments)
                    
                    tool_call_message = format_tool_call(tool_name, arguments)
                    st.session_state.messages.append({"role": "assistant", "content": tool_call_message})
                    with st.chat_message("assistant"):
                        st.info(tool_call_message)
                    
                    # 保存工具调用信息用于后续发送给模型
                    tool_messages.append({
                        "role": "assistant",
                        "content": None,
                        "function_call": {
                            "name": tool_name,
                            "arguments": arguments_str
                        }
                    })
                    
                    # 执行工具调用
                    if st.session_state.is_connected:
                        with st.chat_message("system"):
                            st.info(f"⚙️ 正在执行工具 {tool_name}...")
                        tool_result = await execute_tool_call(tool_name, arguments)
                    else:
                        tool_result = {"result": "服务器未连接，无法执行工具调用"}
                    
                    # 格式化工具调用结果
                    result_message = format_tool_result(tool_result)
                    
                    # 显示工具调用结果
                    with st.chat_message("system"):
                        st.success(f"✅ 工具 {tool_name} 执行完成")
                        with st.expander("查看执行结果", expanded=False):
                            st.json(tool_result)
                    
                    # 保存工具调用结果用于后续发送给模型
                    tool_messages.append({
                        "role": "function",
                        "name": tool_name,
                        "content": json.dumps(tool_result, ensure_ascii=False)
                    })
                
                # 显示工具调用完成
                with st.chat_message("system"):
                    st.success("✨ 所有工具调用已完成，正在生成最终回复...")
                
                # 构建新的对话历史，包含原始提问和工具调用结果
                messages_for_model = [
                    {"role": "user", "content": original_prompt}  # 使用原始提问
                ]
                
                # 添加初始助手回复（如果有）
                if assistant_message:
                    messages_for_model.append({
                        "role": "assistant",
                        "content": assistant_message
                    })
                
                # 添加工具调用消息
                messages_for_model.extend(tool_messages)
                
                # 调用模型获取最终回复
                with st.chat_message("system"):
                    st.info("🤖 正在生成最终回复...")
                final_response = call_tongyi_api_with_messages(messages_for_model)
                
                if "error" not in final_response:
                    final_output = final_response.get("output", {})
                    if "choices" in final_output:
                        final_choice = final_output["choices"][0]
                        final_message = final_choice.get("message", {})
                        final_content = final_message.get("content", "")
                        
                        if final_content:
                            # 如果使用了知识库，在最终回复前添加标记
                            if knowledge_text:
                                final_content = f"📚 基于知识库内容回答：\n\n{final_content}"
                            st.session_state.messages.append({"role": "assistant", "content": final_content})
                            with st.chat_message("assistant"):
                                # 添加打字机效果 - 将最终回复内容逐步显示
                                final_placeholder = st.empty()
                                # 将文本分割成更小的块(按句子或段落)
                                text_chunks = re.split(r'([。！？.!?\n]+)', final_content)
                                displayed_text = ""
                                
                                for i in range(0, len(text_chunks), 2):
                                    # 添加当前块
                                    if i < len(text_chunks):
                                        displayed_text += text_chunks[i]
                                    # 添加分隔符(如果有的话)
                                    if i+1 < len(text_chunks):
                                        displayed_text += text_chunks[i+1]
                                    # 更新显示
                                    final_placeholder.markdown(displayed_text)
                                    # 添加延迟
                                    await asyncio.sleep(0.1)  # 可以调整延迟时间控制"打字"速度
                else:
                    error_message = f"获取最终回复失败: {final_response['error']}"
                    st.session_state.messages.append({"role": "assistant", "content": error_message})
                    with st.chat_message("assistant"):
                        st.error(error_message)
        
        # 在回答完成后显示耗时信息
        with st.chat_message("system"):
            st.markdown(f"""
            <div class='system-message' style='background-color: #f8f9fa;'>
                ⏱️ 回答完成！
                - 总耗时：{total_time:.2f} 秒
                - 完成时间：{time.strftime('%Y-%m-%d %H:%M:%S')}
            </div>
            """, unsafe_allow_html=True)
    
    except Exception as e:
        error_message = f"处理用户查询时发生错误: {str(e)}"
        logger.error(error_message)
        logger.exception("详细错误栈:")
        st.session_state.messages.append({"role": "assistant", "content": error_message})
        with st.chat_message("assistant"):
            st.error(error_message)
        
        # 即使发生错误也显示耗时
        total_time = time.time() - start_time
        with st.chat_message("system"):
            st.markdown(f"""
            <div style='padding: 10px; border-radius: 8px; background-color: #ffebee; margin: 10px 0;'>
                ⚠️ 处理出错！
                - 总耗时：{total_time:.2f} 秒
                - 完成时间：{time.strftime('%Y-%m-%d %H:%M:%S')}
            </div>
            """, unsafe_allow_html=True)

# 侧边栏配置
with st.sidebar:
    # 知识库设置
    st.title("知识库设置")
    use_knowledge_base = st.checkbox("启用知识库", value=st.session_state.use_knowledge_base)
    st.session_state.use_knowledge_base = use_knowledge_base
    
    if use_knowledge_base:
        if not st.session_state.datasets:
            load_datasets()
        
        if st.session_state.datasets:
            dataset_options = {dataset["name"]: dataset["id"] for dataset in st.session_state.datasets}
            selected_dataset_name = st.selectbox(
                "选择知识库",
                options=list(dataset_options.keys()),
                index=list(dataset_options.keys()).index(
                    next((name for name, id in dataset_options.items() if id == st.session_state.selected_dataset), 
                         list(dataset_options.keys())[0])
                )
            )
            st.session_state.selected_dataset = dataset_options[selected_dataset_name]
        else:
            st.warning("未找到可用的知识库")
    else:
        st.info("知识库功能已禁用")

    st.title("服务器配置")
    
    # SSE 服务器配置
    sse_url = st.text_input("SSE 服务器地址", value=st.session_state.sse_url)
    
    # 如果URL改变，更新会话状态
    if sse_url != st.session_state.sse_url:
        st.session_state.sse_url = sse_url
        st.session_state.is_connected = False
        logger.info(f"服务器地址已更改为: {sse_url}")
    
    col1, col2 = st.columns(2)
    
    with col1:
        connect_button = st.button("连接服务器", use_container_width=True)
    
    with col2:
        refresh_button = st.button("刷新工具列表", use_container_width=True, disabled=not st.session_state.is_connected)
    
    disconnect_button = st.button("断开连接", use_container_width=True, disabled=not st.session_state.is_connected)
    
    # 连接服务器
    if connect_button:
        st.session_state.sse_url = sse_url
        connect_to_server()
    
    # 刷新工具列表
    if refresh_button:
        with st.spinner("正在刷新工具列表..."):
            try:
                # 获取工具列表
                tools = asyncio.run(get_tools_from_mcp(st.session_state.sse_url))
                if tools:
                    st.session_state.tools = tools
                    st.success("✅ 工具列表已更新！")
                    logger.info(f"成功更新工具列表: {len(tools)} 个工具")
                else:
                    st.error("❌ 获取工具列表失败")
                    logger.error("获取工具列表失败")
            except Exception as e:
                st.error(f"❌ 刷新工具列表失败: {str(e)}")
                logger.error(f"刷新工具列表失败: {str(e)}")
    
    # 断开连接
    if disconnect_button:
        with st.spinner("正在断开连接..."):
            asyncio.run(safe_disconnect())
            st.session_state.is_connected = False
            st.info("已断开连接")
            logger.info("用户手动断开连接")
    
    # 显示连接状态
    if st.session_state.is_connected:
        st.success("✅ 已连接到服务器")
    else:
        st.warning("⚠️ 未连接到服务器")
    
    # 显示可用工具
    if st.session_state.tools:
        st.subheader(f"可用工具 ({len(st.session_state.tools)})")
        for tool in st.session_state.tools:
            # 适配新的工具格式
            if "function" in tool:
                name = tool["function"]["name"]
                description = tool["function"]["description"]
            else:
                name = tool.get("name", "未命名工具")
                description = tool.get("description", "无描述")
            st.write(f"- {name}: {description}")
    

# 主界面
st.title("AI 智能导购助手")
st.markdown("👋 欢迎使用 AI 智能导购助手！请告诉我您的需求，我会尽力帮助您。")

# 显示示例查询
st.subheader("💡 您可以尝试这些问题:")

# 创建两个容器：一个用于示例按钮，一个用于显示回答
example_buttons = st.container()
response_container = st.container()

# 在第一个容器中显示示例按钮
with example_buttons:
    cols = st.columns(2)
    for i, query in enumerate(EXAMPLE_QUERIES):
        col_idx = i % 2
        with cols[col_idx]:
            if st.button(query, key=f"example_{i}"):
                # 添加用户消息
                st.session_state.messages.append({"role": "user", "content": query})
                # 在完整宽度的容器中显示回答
                with response_container:
                    # 显示用户消息
                    with st.chat_message("user"):
                        st.markdown(query)
                    # 处理用户查询
                    asyncio.run(process_user_query(query))

# 用户输入
prompt = st.chat_input("请输入您的问题")
if prompt:
    # 添加用户消息
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # 处理用户查询
    asyncio.run(process_user_query(prompt))
